#include "asan_avltree.h"

#include "asan_interceptors.h"
#include "asan_internal_defs.h"


static __sanitizer::uptr max(__sanitizer::uptr a, __sanitizer::uptr b) {
  return (a > b) ? a : b;
}

static int height(__sanitizer::AsanChunk *node) {
  if (node == nullptr) return 0;
  return node->height;
}

int __sanitizer::Balance(__sanitizer::AsanChunk *node) {
  if (node == nullptr) return 0;
  return height(node->left) - height(node->right);
}

__sanitizer::AsanChunk *__sanitizer::NewAsanChunk(__sanitizer::uptr beg_addr, __sanitizer::uptr end_addr, int mode) {
  __sanitizer::AsanChunk *node = (__sanitizer::AsanChunk *)REAL(malloc(sizeof(__sanitizer::AsanChunk)));
  node->beg_addr = beg_addr;
  node->end_addr = end_addr;
  node->mode = mode;
  node->height = 1;
  node->left = nullptr;
  node->right = nullptr;
  node->father = nullptr;
  return node;
}

void __sanitizer::UpdateHeight(__sanitizer::AsanChunk *node) {
  node->height = 1 + max(height(node->left), height(node->right));
}

__sanitizer::AsanChunk *__sanitizer::RightRotate(__sanitizer::AsanChunk *node) {
  __sanitizer::AsanChunk *leftChild = node->left;
  __sanitizer::AsanChunk *rightGrandChild = leftChild->right;

  leftChild->right = node;
  node->father = leftChild;
  node->left = rightGrandChild;

  if (rightGrandChild != nullptr) rightGrandChild->father = node;

  __sanitizer::UpdateHeight(node);
  __sanitizer::UpdateHeight(leftChild);

  return leftChild;
}

__sanitizer::AsanChunk *__sanitizer::LeftRotate(__sanitizer::AsanChunk *node) {
  __sanitizer::AsanChunk *rightChild = node->right;
  __sanitizer::AsanChunk *leftGrandChild = rightChild->left;

  rightChild->left = node;
  node->father = rightChild;
  node->right = leftGrandChild;

  if (leftGrandChild != nullptr) leftGrandChild->father = node;

  __sanitizer::UpdateHeight(node);
  __sanitizer::UpdateHeight(rightChild);

  return rightChild;
}

__sanitizer::AsanChunk *__sanitizer::InsertChunk(__sanitizer::AsanChunk *node, __sanitizer::AsanChunk *new_node) {
  if (node == nullptr) return new_node;

  if (__sanitizer::ChunkAtLeft(node, new_node)) {
    node->left = __sanitizer::InsertChunk(node->left, new_node);
    node->left->father = node;
  }
  else if (__sanitizer::ChunkAtRight(node, new_node)) {
    node->right = __sanitizer::InsertChunk(node->right, new_node);
    node->right->father = node;
  }
  else {
    Printf("Insert an existing node (beg_addr=%p).\n", (void*)(new_node->beg_addr));
    Die();
  }

  __sanitizer::UpdateHeight(node);

  int BalanceFactor = __sanitizer::Balance(node);

  if (BalanceFactor > 1 && __sanitizer::ChunkAtLeft(node->left, new_node)) {
    return __sanitizer::RightRotate(node);
  }

  if (BalanceFactor < -1 && __sanitizer::ChunkAtRight(node->right, new_node)) {
    return __sanitizer::LeftRotate(node);
  }

  if (BalanceFactor > 1 && __sanitizer::ChunkAtRight(node->left, new_node)) {
    node->left = __sanitizer::LeftRotate(node->left);
    node->left->father = node;
    return __sanitizer::RightRotate(node);
  }

  if (BalanceFactor < -1 && __sanitizer::ChunkAtLeft(node->right, new_node)) {
    node->right = __sanitizer::RightRotate(node->right);
    node->right->father = node;
    return __sanitizer::LeftRotate(node);
  }

  return node;
}

__sanitizer::AsanChunk *__sanitizer::MinValueChunk(__sanitizer::AsanChunk *node) {
  __sanitizer::AsanChunk *current = node;
  while (current->left != nullptr) current = current->left;
  return current;
}

__sanitizer::AsanChunk *__sanitizer::MaxValueChunk(__sanitizer::AsanChunk *node) {
  __sanitizer::AsanChunk *current = node;
  while (current->right != nullptr) current = current->right;
  return current;
}

__sanitizer::AsanChunk *__sanitizer::DeleteChunk(__sanitizer::AsanChunk *node, __sanitizer::uptr beg_addr) {
  if (node == nullptr) return node;

  if (beg_addr < node->beg_addr) {
    node->left = __sanitizer::DeleteChunk(node->left, beg_addr);
    if (node->left != nullptr) node->left->father = node;
  }

  else if (beg_addr > node->beg_addr) {
    node->right = __sanitizer::DeleteChunk(node->right, beg_addr);
    if (node->right != nullptr) node->right->father = node;
  }
  else {
    if ((node->left == nullptr) || (node->right == nullptr)) {
      __sanitizer::AsanChunk *temp = (node->left != nullptr) ? node->left : node->right;

      if (temp == nullptr) {
        temp = node;
        if (node->father != nullptr) {
          if (node->father->left == node) node->father->left = nullptr;
          else node->father->right = nullptr;
        }
        node = nullptr;
      }
      else {
        __sanitizer::AsanChunk* node_father = node->father;
        *node = *temp;
        if (node->left != nullptr) node->left->father = node;
        if (node->right != nullptr) node->right->father = node;
        node->father = node_father;
      }
      REAL(free(temp));
    }
    else {
      __sanitizer::AsanChunk *temp = __sanitizer::MinValueChunk(node->right);
      node->beg_addr = temp->beg_addr;
      node->end_addr = temp->end_addr;
      node->mode = temp->mode;
      node->right = __sanitizer::DeleteChunk(node->right, temp->beg_addr);
      if (node->right != nullptr)
        node->right->father = node;
    }
  }

  if (node == nullptr) return node;

  __sanitizer::UpdateHeight(node);

  int BalanceFactor = __sanitizer::Balance(node);

  if (BalanceFactor > 1 && __sanitizer::Balance(node->left) >= 0) {
    return __sanitizer::RightRotate(node);
  }

  if (BalanceFactor > 1 && __sanitizer::Balance(node->left) < 0) {
    node->left = __sanitizer::LeftRotate(node->left);
    node->left->father = node;
    return __sanitizer::RightRotate(node);
  }

  if (BalanceFactor < -1 && __sanitizer::Balance(node->right) <= 0) {
    return __sanitizer::LeftRotate(node);
  }

  if (BalanceFactor < -1 && __sanitizer::Balance(node->right) > 0) {
    node->right = __sanitizer::RightRotate(node->right);
    node->right->father = node;
    return __sanitizer::LeftRotate(node);
  }

  return node;
}
